# -*- coding: utf-8 -*-

import sys
import numpy as np
from tqdm import tqdm
import re
from utils.alphabet import Alphabet

# from transformers.tokenization_bert import BertTokenizer
from transformers import BertTokenizer
from loguru import logger

NULLKEY = "-null-" # 用于填充的特殊标记，表示空值或无效值

#模式定义：这些常量用于定义批次的类型和处理方式

# 每个批次包含固定大小的实例（训练和测试通用）
MODE_INSTANCE = (
    0  # One batch contains batch_size instances. [for both training and testing]
)
# 每个批次包含多个袋子（每个袋子内的实例共享相同的实体对，通常用于测试）
MODE_ENTPAIR_BAG = 1  # One batch contains batch_size bags, instances in which have the same entity pair [for testing].
# 每个批次包含多个袋子（每个袋子内的实例共享相同的关系事实，通常用于训练）
MODE_RELFACT_BAG = 2  # One batch contains batch size bags, instances in which have the same relation fact. [for training].
# 每个批次包含多个实例，每个实例可能有多个标签（用于多标签分类的测试）
MODE_INSTANCE_MULTI_LABEL = 3  # One batch contains batch_size instances, and each answer is multi-label (there might be more than 1 correct answer) [for testing].

#重复的代码
# 规范化单词，将其中的数字替换为"0"
def normalize_word(word):
    """
    所谓NORMALIZED WORD 就是将数字抹掉
    """
    new_word = ""
    for char in word:
        if char.isdigit():
            new_word += "0"
        else:
            new_word += char
    return new_word


# transform string to list
# '第<N>天' -> ['第','<N>','天']
#重复的代码
# 将字符串转换为字符列表，处理特殊的符号分割
def str2list(str, spec=None):
    str = re.sub("\s+", "", str)
    res = []
    i = 0
    while i < len(str):
        match = False
        if spec:
            for sp in spec:
                splen = len(sp)
                if i + splen <= len(str) and str[i : i + splen] == sp:
                    res.append(sp)
                    i += splen
                    match = True
                    break
        if not match:
            res.append(str[i])
            i += 1
    return res


def load_data(input_file, mode):
    """
   从输入文件中加载数据并进行预处理,字典返回数据和标签频率
   :param input_file: 输入文件路径
   :param mode: 模式，影响数据加载方式
   :return: 返回处理后的数据和标签频率
   """
    ori_data = [] # 存储原始数据
    freq = dict()  # 存储每个标签的出现频率

    for line in open(input_file, "r", encoding="utf-8"):
        line = line.strip().split("\t")  #按制表符分割每行数据
        if len(line) == 5:# 解析5个字段
            sentid, ent1, ent2, label, sent = line
        elif len(line) == 4:# 解析4个字段
            ent1, ent2, label, sent = line
        else: # 如果数据格式不正确，跳过该行
            continue
        # 检查实体是否出现在句子中
        if ent1 not in sent:
            logger.debug(ent1 + " not found in " + sent)
            continue
        if ent2 not in sent:
            logger.debug(ent2 + " not found in " + sent)
            continue
        # 统计标签频率
        if label not in freq:
            freq[label] = 1
        else:
            freq[label] += 1
        # 将每条数据转换为字典形式
        ori_data.append(
            {"head": ent1, "tail": ent2, "relation": label, "sentence": sent}
        )

    return ori_data, freq


def read_instance_with_gaz(
    input_file,
    # ----- 字符级别字典和词级别字典对于汉字都是一样的 这三个都是Alphabet类
    char_alphabet,
    word_alphabet,
    biword_alphabet,
    # -----
    gaz,#gazetteer对象
    gaz_alphabet,#Alphabet类
    gaz_count,#词典频率dict
    # ----
    label_alphabet,#label的Alphabet类
    number_normalized,
    max_sent_length,
    char_padding_size=-1,
    char_padding_symbol="</pad>",
    mode=1,
):
    """
    读取带有实体词典（gazetteer）信息的实例，并进行特征处理
    :param input_file: 输入文件路径
    :param char_alphabet: 字符级别的字典
    :param word_alphabet: 词级别的字典
    :param biword_alphabet: 双词级别的字典
    :param gaz: 实体词典对象
    :param gaz_alphabet: 实体词典的字典
    :param gaz_count: 实体词典中的实体计数
    :param label_alphabet: 标签字典
    :param number_normalized: 是否进行数字规范化
    :param max_sent_length: 最大句子长度
    :param char_padding_size: 字符级别的填充大小，-1不填充
    :param char_padding_symbol: 字符级别的填充符号
    :param mode: 模式，控制数据加载方式
    :return: 返回文本和索引的处理结果
    """

    # TODO

    from model.config_layer import config_semantic

    bert_path = config_semantic["bert_path"]
    # bert_path = "hf_models/chinese-bert-wwm-ext"
    # bert_path = "hf_models/chinese-lert-base"

    tokenizer = BertTokenizer.from_pretrained(bert_path, do_lower_case=True)

    instence_texts = [] # 存储文本特征
    instence_Ids = [] # 存储索引特征

    ori_data, freq = load_data(input_file, mode=mode)# 加载原始数据
#ori_data: {‘head':'','relation':'','sentence':'','tail':''}  freq: {'label':count}
    UNK_id = gaz_alphabet.get_index(gaz_alphabet.UNKNOWN)# 获取未知词的索引
    for data in ori_data:

        words, word_Ids = [], [] # 存储词和词索引
        biwords, biword_Ids = [], [] # 存储双词和双词索引
        chars, char_Ids = [], [] # 存储字符和字符索引
        labels, label_Ids = [], [] # 存储标签和标签索引
        # 处理标签部分，关系类型
        labels.append(data["relation"])# 获取样本的关系类型 一维list
        label_Ids.append(label_alphabet.get_index(data["relation"]))# 将标签转换为索引 一维list
        # 获取实体1、实体2和句子
        ent1, ent2, sent = data["head"], data["tail"], data["sentence"]
        ent1, ent2, sent = str2list(ent1), str2list(ent2), str2list(sent)  # 将实体和句子转为列表 一维list

        labels = [data["relation"]]# 获取样本的关系类型 一维list,append失效
        # 规范化处理
        ent1 = "".join(ent1)#变成字符串
        ent2 = "".join(ent2)#变成字符串
        # 如果句子末尾没有句号且句子长度小于最大长度，则添加句号
        if sent[-1] not in ["。"] and len(sent) < max_sent_length:
            sent.append("。")
        #ent1字符串，ent2字符串，sent句子原来list
        sent = [ent1, "[SEP]", ent2, "[SEP]"] + sent  # 将实体插入句子中，实体之间用 [SEP] example: 分隔 ['幽兰', '[SEP]', '山谷', '[SEP]', '“', '幽', '兰', '在', '山', '谷', '，', '本', '自', '无', '人', '识', '。']
        # 处理词、双词、字符等特征
        for widx, word in enumerate(sent):
            if number_normalized:
                word = normalize_word(word) # 如果启用规范化，替换数字为0
            if widx < len(sent) - 1 and len(sent) > 2:
                biword = word + sent[widx + 1]  # 双词
            else:
                biword = word + NULLKEY # 处理结尾的特殊标记

            # words and bi-gram (not used)
            biwords.append(biword)#'幽兰[SEP]'
            biword_Ids.append(biword_alphabet.get_index(biword))#如果词典中没有会添加

            words.append(word)#['幽兰']
            word_Ids.append(word_alphabet.get_index(word))

            # character features (not uesd)
            # 处理字符级特征（暂时没有用）
            char_list = []#['幽', '兰']
            char_Id = []
            for char in word:
                char_list.append(char)

            if char_padding_size > 0: # 如果需要字符填充，填充pad字符到指定长度
                char_number = len(char_list)
                if char_number < char_padding_size:
                    char_list = char_list + [char_padding_symbol] * (
                        char_padding_size - char_number
                    )
                assert len(char_list) == char_padding_size
            for char in char_list:
                char_Id.append(char_alphabet.get_index(char))

            chars.append(char_list)#[['幽', '兰']]
            char_Ids.append(char_Id)
        # 检查句子长度是否符合最大长度要求
        if ((max_sent_length < 0) or (len(words) < max_sent_length)) and (
            len(words) > 0
        ):
            # 继续处理实体匹配和gazetteer（词典）相关的信息
            gaz_Ids = [] #存储实体匹配的索引
            layergazmasks = []  # 存储gazetteer的mask
            gazchar_masks = [] # 存储字符级别的mask
            w_length = len(words)#words ['幽兰', '[SEP]', '山谷', '[SEP]', '“', '幽', '兰', '在', '山', '谷', '，', '本', '自', '无', '人', '识', '。']
            #维度 [w_length,4,未知长度list]
            gazs = [[[] for i in range(4)] for _ in range(w_length)]# 匹配的实体词典
            # gazs:[c1,c2,...,cn]  ci:[B,M,E,S]  B/M/E/S :[w_id1,w_id2,...]  None:0
            gazs_count = [[[] for i in range(4)] for _ in range(w_length)]# 匹配实体的数量
            gaz_char_Id = [[[] for i in range(4)] for _ in range(w_length)] # 存储字符级别的实体信息
            ## gazs:[c1,c2,...,cn]  ci:[B,M,E,S]  B/M/E/S :[[w1c1,w1c2,...],[],...]

            max_gazlist = 0 # 最大实体数目
            max_gazcharlen = 0 # 最大字符长度
            for idx in range(w_length):
                # 匹配gazetteer中的实体
                matched_list = gaz.enumerateMatchList(words[idx:])# 从gazetteer中获取匹配的实体列表 传入后缀list
                matched_length = [len(a) for a in matched_list]# 获取匹配的实体长度
                matched_Id = [gaz_alphabet.get_index(entity) for entity in matched_list]# 获取实体的索引
                #example: ['我'] [1] [114]
                if matched_length:#计算最大词典长度
                    max_gazcharlen = max(max(matched_length), max_gazcharlen)

                for w in range(len(matched_Id)):
                    gaz_chars = []
                    g = matched_list[w] #matched_Id对应的单词 example:我
                    for c in g:  # 将匹配的gaz拆开是有何用意？ # 将匹配的gazetteer实体拆分为字符索引
                        gaz_chars.append(word_alphabet.get_index(c))
                    #BMES
                    if matched_length[w] == 1:  ## Single # 单一实体
                        gazs[idx][3].append(matched_Id[w]) # 单一实体id存入'S'位置
                        gazs_count[idx][3].append(1)
                        gaz_char_Id[idx][3].append(gaz_chars)
                    else:
                        gazs[idx][0].append(matched_Id[w])  ## Begin# 起始位置存入'B'
                        gazs_count[idx][0].append(gaz_count[matched_Id[w]])
                        gaz_char_Id[idx][0].append(gaz_chars)#字符id list

                        wlen = matched_length[w]
                        gazs[idx + wlen - 1][2].append(matched_Id[w])  ## End # 结束位置存入'E'
                        gazs_count[idx + wlen - 1][2].append(gaz_count[matched_Id[w]])

                        gaz_char_Id[idx + wlen - 1][2].append(gaz_chars)
                        for l in range(wlen - 2):
                            gazs[idx + l + 1][1].append(matched_Id[w])  ## Middle # 中间位置存入'M'
                            gazs_count[idx + l + 1][1].append(gaz_count[matched_Id[w]])
                            gaz_char_Id[idx + l + 1][1].append(gaz_chars)
                # 补齐实体的mask信息
                for label in range(4):
                    if not gazs[idx][label]:
                        # 为何要在所有的 BMES 匹配后面添加0 否则填字典id
                        gazs[idx][label].append(0)# 如果没有匹配到实体，填充0
                        gazs_count[idx][label].append(1)
                        gaz_char_Id[idx][label].append([0])

                    max_gazlist = max(len(gazs[idx][label]), max_gazlist)

                matched_Id = [
                    gaz_alphabet.get_index(entity) for entity in matched_list
                ]  # 词号 # 将实体转为索引 [114]
                if matched_Id:
                    gaz_Ids.append([matched_Id, matched_length]) # 存储匹配的实体id,实体长度
                else:
                    gaz_Ids.append([])# 如果没有匹配到实体，存储空数据
            ## batch_size = 1
            # 处理字符级别的mask
            for idx in range(w_length):
                """
                针对每一个char匹配的gaz 和 gaz 对应的 gaz_char
                """
                gazmask = []
                gazcharmask = []

                for label in range(4):
                    label_len = len(gazs[idx][label])
                    count_set = set(gazs_count[idx][label])
                    if len(count_set) == 1 and 0 in count_set:
                        gazs_count[idx][label] = [1] * label_len

                    mask = label_len * [0]
                    mask += (max_gazlist - label_len) * [1]  # 1 为 mask 的值

                    gazs[idx][label] += (max_gazlist - label_len) * [0]  ## padding
                    gazs_count[idx][label] += (max_gazlist - label_len) * [
                        0
                    ]  ## padding

                    char_mask = []
                    for g in range(len(gaz_char_Id[idx][label])):
                        glen = len(gaz_char_Id[idx][label][g])
                        charmask = glen * [0]
                        charmask += (max_gazcharlen - glen) * [1]
                        char_mask.append(charmask)
                        gaz_char_Id[idx][label][g] += (max_gazcharlen - glen) * [0]
                    gaz_char_Id[idx][label] += (max_gazlist - label_len) * [
                        [0 for i in range(max_gazcharlen)]
                    ]
                    char_mask += (max_gazlist - label_len) * [
                        [1 for i in range(max_gazcharlen)]
                    ]

                    gazmask.append(mask)
                    gazcharmask.append(char_mask)

                layergazmasks.append(gazmask)
                gazchar_masks.append(gazcharmask)
            # 使用BERT tokenizer对文本进行标记化
            texts = ["[CLS]"] + words + ["[SEP]"]
            bert_text_ids = tokenizer.convert_tokens_to_ids(texts)

            instence_texts.append([words, biwords, chars, gazs, labels])# 存储文本特征
            instence_Ids.append(
                [
                    word_Ids,
                    biword_Ids,
                    char_Ids,
                    gaz_Ids,
                    label_Ids,
                    # --------------
                    gazs,
                    gazs_count,
                    gaz_char_Id,
                    # -----------
                    layergazmasks,
                    gazchar_masks,
                    # --------
                    bert_text_ids,
                ]
            )# 存储索引特征
        # breakpoint()

    return instence_texts, instence_Ids


def build_pretrain_embedding(embedding_path, word_alphabet, embedd_dim=100, norm=True):
    logger.info("Build from {}...".format(embedding_path))
    # 创建一个空字典，用于存储预训练词向量
    embedd_dict = dict()
    # 如果指定了预训练词嵌入路径，则加载预训练嵌入
    if embedding_path != None:
        embedd_dict, embedd_dim = load_pretrain_emb(embedding_path)
    # 使用一个小的尺度因子来初始化嵌入，避免初始值过大
    scale = np.sqrt(3.0 / embedd_dim)
    # 创建一个大小为(word_alphabet.size(), embedd_dim)的空numpy数组
    # 用于存储每个词的嵌入向量
    pretrain_emb = np.empty([word_alphabet.size(), embedd_dim])
    # 初始化一些计数器，用于统计匹配结果
    perfect_match = 0 # 完全匹配的词数
    case_match = 0  # 只大小写匹配的词数
    not_match = 0 # 无匹配的词数
    # 对于词汇表中的第一个词，随机初始化一个嵌入向量（通常是用于PAD标记或其他特殊标记）
    pretrain_emb[0, :] = np.random.uniform(-scale, scale, [1, embedd_dim])
    # 遍历词汇表中的每个词及其对应的索引
    for word, index in tqdm(word_alphabet.instance2index.items()):
        # 如果词汇表中的词在预训练词嵌入字典中找到匹配
        if word in embedd_dict:
            # 如果启用了归一化（norm=True），则将词向量归一化为单位向量
            if norm:
                pretrain_emb[index, :] = norm2one(embedd_dict[word])
            else:
                pretrain_emb[index, :] = embedd_dict[word]
            # 记录完全匹配的词数量
            perfect_match += 1
        # 如果词汇表中的词的小写版本在预训练词嵌入字典中找到匹配
        elif word.lower() in embedd_dict:
            # 如果启用了归一化，则归一化该词的小写嵌入
            if norm:
                pretrain_emb[index, :] = norm2one(embedd_dict[word.lower()])
            else:
                pretrain_emb[index, :] = embedd_dict[word.lower()]
            # 记录大小写匹配的词数量
            case_match += 1
        else:
            # 如果没有找到匹配的词，则随机初始化该词的嵌入
            pretrain_emb[index, :] = np.random.uniform(-scale, scale, [1, embedd_dim])
            not_match += 1 # 记录未匹配的词数量、
    # 获取预训练词向量字典的大小
    pretrained_size = len(embedd_dict)
    # 打印出匹配结果的统计信息
    print(
        "Embedding:\n     pretrain word:%s, prefect match:%s, case_match:%s, oov:%s, oov%%:%s"
        % (
            pretrained_size,
            perfect_match,
            case_match,
            not_match,
            (not_match + 0.0) / word_alphabet.size(),
        )
    )
    # 返回最终的词嵌入矩阵和嵌入维度
    return pretrain_emb, embedd_dim


def norm2one(vec):
    # 计算向量 vec 中所有元素的平方和，然后取平方根，得到向量的 L2 范数（即模长）
    root_sum_square = np.sqrt(np.sum(np.square(vec)))
    # 将向量 vec 中的每个元素除以 L2 范数，实现归一化操作
    return vec / root_sum_square


def load_pretrain_emb(embedding_path):
    embedd_dim = -1 # 初始化嵌入维度为-1，表示尚未确定
    embedd_dict = dict()  # 创建一个空字典，用于存储加载的词与词向量的对应关系
    # 打开预训练词向量文件，按行读取文件
    with open(embedding_path, "r", encoding="utf-8") as file:
        for line in tqdm(file):
            line = line.strip()
            if len(line) == 0 or len(line) == 2:
                continue
            tokens = line.split()
            if len(tokens) == 2:
                fmt = "Total {} word, dim  {}\n".format(tokens[0], tokens[1])
                print(fmt)
                continue
            if embedd_dim < 0:
                embedd_dim = len(tokens) - 1
            else:
                assert embedd_dim + 1 == len(tokens)

            embedd = np.empty([1, embedd_dim])  # 创建一个空的NumPy数组用于存储该词的嵌入向量
            embedd[:] = tokens[1:]  # 将当前行中的嵌入向量（即从第二项开始的所有值）赋值给该词的嵌入向量
            embedd_dict[tokens[0]] = embedd  # 将词和其对应的嵌入向量添加到字典中
    # 返回加载的嵌入字典和嵌入的维度
    return embedd_dict, embedd_dim
